#pragma once

#include "client/client.hpp"
#include "websocket.hpp"

#include "TinySHA1.hpp"

/**
 * @brief a websocket client template based on an abstract socket client
 *
 * @tparam client the socket client implementation
 */
template <class client_t = client>
class websocket_client {

    static constexpr auto WEBSOCKET_SERVER_MAGIC_KEY = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    static constexpr auto HEARTBEAT_INTERVAL = std::chrono::seconds(5);
    static constexpr auto KEEPALIVE_INTERVAL = HEARTBEAT_INTERVAL * 3;

public:
    using recv_hook_t = std::function<bool(websocket_frame)>;

    websocket_client()
        : _buffer([this](websocket_frame frame) { return _handle_websocket_frame(std::move(frame)); }, false)
        , _mt19937(_random_device())
        , _int_dist(0, std::numeric_limits<uint8_t>::max()) {
        _recv_hook = [](websocket_frame frame) {
            linfo("received websocket data from server: \n" << frame.description());
            return true;
        };
        _client.register_recv_hook([this](std::vector<char> buf) {
            auto res = _handle_recv(std::move(buf));
            if (!res) {
                _state = websocket_state::CLOSED;
            }
            return res;
        });
        _client.register_close_hook([this]() {
            if (std::this_thread::get_id() != _daemon_thread.get_id()) {
                {
                    std::lock_guard<std::mutex> lock(_daemon_mutex);
                    _daemon_exit = true;
                    _daemon_cv.notify_one();
                }
                _daemon_thread.join();
            }
        });
    }

    void register_recv_hook(recv_hook_t recv_hook) {
        _recv_hook = std::move(recv_hook);
    }

    [[nodiscard]] bool send_data(std::vector<char> buf, websocket_frame_code code) {
        // #ifdef FRAGMENT_DATA
        constexpr size_t FRAGMENT_SIZE = 1 << 20;
        if (buf.size() < FRAGMENT_SIZE) {
            return send_frame(std::vector(buf.data(), buf.data() + buf.size()), code, true);
        }
        const int total_size = static_cast<int>(buf.size());
        int rear = 0;
        while (total_size - rear > FRAGMENT_SIZE) {
            std::vector cur(buf.data() + rear, buf.data() + rear + FRAGMENT_SIZE);
            rear += FRAGMENT_SIZE;
            if (!send_frame(std::move(cur), code, false)) {
                return false;
            }
            code = CONTINUATION_FRAME;
        }
        return send_frame(std::vector(buf.data() + rear, buf.data() + total_size), code, true);
        // #else
        // return send_frame(std::move(buf), code, true);
        // #endif
    }

    void close(uint16_t close_cause = 1000) {
        websocket_frame frame;
        frame.opcode = CONNECTION_CLOSE;
        frame.masked = true;
        frame.buffer = {static_cast<char>(close_cause / 256), static_cast<char>(close_cause % 256)};
        frame.length = frame.buffer.size();
        _state = websocket_state::CLOSING;
        if (!_client.send(frame.dump())) {
            lerror("cannot send close websocket frame.");
        }
        linfo("close frame sent to server.");
    }

    [[nodiscard]] bool alive() const {
        return _client.alive();
    }

    [[nodiscard]] bool connect(const std::string& address, u_short port) {
        if (!_client.connect(address, port)) {
            linfo("connect to server " << address << ":" << port << " failed.");
            return false;
        }

        std::string key(16, '\0');
        for (auto& c: key) {
            c = random_uint8();
        }
        key = base64_encode(key);
        handshake_request request;
        request.method = "GET";
        request.uri = "/";
        request.version = "HTTP/1.1";
        request.message.key_values["Host"] = {"127.0.0.1"};
        request.message.key_values["Origin"] = {"127.0.0.1"};
        request.message.key_values["Upgrade"] = {"websocket"};
        request.message.key_values["Connection"] = {"Upgrade"};
        request.message.key_values["Sec-WebSocket-Key"] = {key};
        request.message.key_values["Sec-WebSocket-Version"] = {"13"};

        _expected_accept = sha1_then_base64(key + WEBSOCKET_SERVER_MAGIC_KEY);
        auto v = request.dump();
        if (!_client.send(std::vector(v.data(), v.data() + v.length()))) {
            _client.close();
            _state = websocket_state::CLOSED;
            return false;
        }
        _state = websocket_state::CONNECTING;
        return true;
    }

    [[nodiscard]] bool send_frame(std::vector<char> buf, websocket_frame_code code, bool FIN = true) {
        if (_state != websocket_state::OPEN || !_client.alive()) {
            return false;
        }
        websocket_frame frame;
        frame.opcode = code;
        frame.FIN = FIN;
        frame.masked = true;

        frame.buffer = std::move(buf);
        frame.length = frame.buffer.size();

        frame.mask_key[0] = random_uint8();
        frame.mask_key[1] = random_uint8();
        frame.mask_key[2] = random_uint8();
        frame.mask_key[3] = random_uint8();

        for (size_t i = 0; i < frame.buffer.size(); ++i) {
            frame.buffer[i] ^= frame.mask_key[i % 4];
        }

        return _client.send(frame.dump());
    }

private:
    bool _send_pong(std::vector<char> buf) {
        websocket_frame frame;
        frame.opcode = PONG;
        frame.length = buf.size();
        frame.FIN = true;
        frame.buffer = std::move(buf);
        auto dumped = frame.dump();
        return _client.send(std::move(dumped));
    }

    bool _handle_websocket_frame(websocket_frame frame) {
        switch (frame.opcode) {
        case PING: {
            linfo("received ping from server.");
            auto buf = frame.buffer;
            return _recv_hook(std::move(frame)) && _send_pong(std::move(buf));
        }
        case PONG: {
            int64_t ti = 0, base = 1;
            for (size_t i = 0; i < frame.buffer.size(); i++) {
                ti += base * static_cast<uint8_t>(frame.buffer[i]);
                base *= 256;
            }
            double duration = (now_in_i64() - ti) / 1000.0;
            linfo("received pong from " << duration << " seconds ago");
            return _recv_hook(std::move(frame));
        }
        case CONNECTION_CLOSE: {
            if (_state == websocket_state::CLOSING) {
                linfo("websocket server replied close request");
            } else {
                linfo("websocket close request from server");
                _recv_hook(std::move(frame));
                if (send_frame({}, CONNECTION_CLOSE)) {
                    lwarn("send close frame failed.");
                }
            }
            return false;
        }
        default:;
        }
        return _recv_hook(std::move(frame));
    }

    bool _handle_recv(std::vector<char> buf) {
        if (buf.empty() || _state == websocket_state::CLOSED) return true;
        if (_state == websocket_state::CONNECTING) {
            auto result = handshake_response::parse_handshake_response(buf);
            if (result.index() == 0) { // std::string
                lerror("handshake response error: " << std::get<std::string>(result));
                return false;
            }
            auto response = std::get<handshake_response>(result);
            if (!response.is_valid_response(_expected_accept)) {
                lerror("invalid handshake response");
                return false;
            }
            _state = websocket_state::OPEN;
            linfo("received valid handshake response, handshake success.");

            _daemon_exit = false;
            _daemon_thread = std::thread([this]() {
                while (!_daemon_exit && _client.alive()) {
                    std::unique_lock<std::mutex> uni_lock(_daemon_mutex);
                    _daemon_cv.wait_for(uni_lock, HEARTBEAT_INTERVAL, [&] { return _daemon_exit.load(); });
                    if (_daemon_exit) {
                        break;
                    }
                    auto now = now_in_i64();
                    std::vector<char> buf(8);
                    for (int i = 0; i < 8; i++) {
                        buf[i] = now % 256;
                        now /= 256;
                    }
                    if (!send_frame(std::move(buf), PING)) {
                        lerror("websocket server disconnected, closing");
                        _client.close();
                        _state = websocket_state::CLOSED;
                        break;
                    }
                    auto diff = std::chrono::duration_cast<std::chrono::milliseconds>(
                        _last_receive_time - std::chrono::system_clock::now());
                    if (diff > KEEPALIVE_INTERVAL) {
                        lerror("websocket server timeout, closing");
                        _client.close();
                        _state = websocket_state::CLOSED;
                        break;
                    }
                }
                linfo("daemon thread exit.");
            });
            return true;
        }
        return _buffer.receive_data(std::move(buf));
    }

    [[nodiscard]] uint8_t random_uint8() {
        return static_cast<uint8_t>(_int_dist(_mt19937));
    }

    std::chrono::time_point<std::chrono::system_clock> _last_receive_time;

    std::random_device _random_device;
    std::mt19937 _mt19937;
    std::uniform_int_distribution<uint32_t> _int_dist;

    std::atomic<bool> _daemon_exit{false};
    std::thread _daemon_thread;

    std::mutex _daemon_mutex;
    std::condition_variable _daemon_cv;

    std::string _expected_accept;
    websocket_state _state{websocket_state::CLOSED};
    websocket_receive_buffer _buffer;
    recv_hook_t _recv_hook{nullptr};
    client_t _client;
};

using websocket_secure_client = websocket_client<secure_client>;